#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import cv2
import struct
import multiprocessing
import time
from collections import OrderedDict
import numpy as np
# import jpeg4py as jpeg

base_path = r'/rootfs/media/kasim/DataSet/RecordEx'
out_path = r'/rootfs/media/kasim/DataSet/FaceCropImage'
video_list_file = r'/rootfs/media/kasim/DataSet/RecordEx/face_video_list_bbox_count.attr.txt'

PROCESS_COUNT = 16

SLEEP_TIME = 1
SHOW = False
SAVE_IMAGE = True
# SHOW = True
# SAVE_IMAGE = False

# VERY_GOOD = 0
# GOOD = 1
# BAD = 2
# VERY_BAD = 3
filter_quality = 3
filter_quality_thr = 0.25

min_pitch = -30
max_pitch = 30
min_roll = -45
max_roll = 45
min_yaw = -30
max_yaw = 30

min_angle_score = 0.2
max_angle_diff = 5

max_size = 448


def resize_image(bbox_image, scale_size):
    w = bbox_image.shape[1]
    h = bbox_image.shape[0]
    size = max(w, h)
    scale = scale_size / size
    if scale < 1.0:
        w = int(round(w * scale))
        h = int(round(h * scale))
        bbox_image = cv2.resize(bbox_image, (w, h))
    return bbox_image


def crop_proc(file_queue, out_queue, id):
    try:
        win_name = 'process{}'.format(id)

        is_break = False
        while True:
            # file_info = file_queue.get(timeout=5)
            file_info = file_queue.get(block=True)
            if file_info is None:
                break

            file_name, attributes_dict = file_info

            img_path = os.path.join(base_path, file_name)
            if not os.path.exists(img_path):
                print(img_path, 'no exists')
                continue

            cap = cv2.VideoCapture(img_path)

            image_dir = os.path.splitext(file_name)[0]
            image_path = os.path.join(out_path, image_dir)
            if not os.path.exists(image_path):
                os.makedirs(image_path)

            frame_id = 0
            while True:
                grabbed, image_bgr = cap.read()

                if not grabbed:
                    break

                width, height = image_bgr.shape[1], image_bgr.shape[0]

                attributes = attributes_dict.get(frame_id, None)
                if attributes is not None:
                    bboxes = []
                    for attribute in attributes:
                        quality = attribute['quality']
                        quality = np.array(quality, dtype=np.float32)
                        max_idx = np.argmax(quality)
                        if (max_idx == filter_quality) and (quality[max_idx] >= filter_quality_thr):
                            continue

                        pitch = attribute['pitch'][0]
                        yaw = attribute['yaw'][0]
                        roll = attribute['roll'][0]

                        if pitch < min_pitch or max_pitch < pitch or roll < min_roll or max_roll < roll or yaw < min_yaw or max_yaw < yaw:
                            continue

                        bin_pitch = attribute['bin_pitch']
                        bin_yaw = attribute['bin_yaw']
                        bin_roll = attribute['bin_roll']

                        bin_pitch_score = float(bin_pitch[1])
                        bin_pitch = float(bin_pitch[0])
                        bin_roll_score = float(bin_roll[1])
                        bin_roll = float(bin_roll[0])
                        bin_yaw_score = float(bin_yaw[1])
                        bin_yaw = float(bin_yaw[0])

                        pitch_diff = abs(bin_pitch - pitch)
                        roll_diff = abs(bin_roll - roll)
                        yaw_diff = abs(bin_yaw - yaw)
                        if max_angle_diff < pitch_diff or max_angle_diff < roll_diff or max_angle_diff < yaw_diff:
                            continue

                        if bin_pitch_score <= min_angle_score or bin_roll_score <= min_angle_score or bin_yaw_score <= min_angle_score:
                            continue

                        bbox = attribute['bboxes']
                        xmin, ymin, xmax, ymax = list(map(int, map(round, bbox)))
                        # xmin = max(xmin, 0)
                        # ymin = max(ymin, 0)
                        # xmax = min(xmax, width)
                        # ymax = min(ymax, height)
                        bboxes.append([xmin, ymin, xmax, ymax])

                    if len(bboxes) > 0:
                        if SHOW:
                            for bbox in bboxes:
                                cv2.rectangle(image_bgr, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), thickness=2)

                        if SAVE_IMAGE:
                            image = image_bgr
                            for i, mbbox in enumerate(bboxes):
                                image_name = 'F{:04d}B{:02d}.jpg'.format(frame_id, i)

                                out_image_path = os.path.join(image_path, image_name)
                                bbox_image = image[mbbox[1]:mbbox[3], mbbox[0]:mbbox[2]]
                                bbox_image = resize_image(bbox_image, max_size)
                                cv2.imwrite(out_image_path, bbox_image)

                frame_id += 1
                if SHOW:
                    cv2.imshow(win_name, image_bgr)

                    if 0 == SLEEP_TIME:
                        k = cv2.waitKey()
                    else:
                        k = cv2.waitKey(SLEEP_TIME)
                    if (k == 27) or (k == ord('q')) or (k == ord('Q')):  # Esc key to stop
                        is_break = True
                        break

            # os.system('chmod a+wr {}'.format(bbox_file_path))
            while out_queue.qsize() > 100:
                time.sleep(0.01)
            out_queue.put((file_name, False, id))
            if is_break:
                break
    except Exception as e:
        print('Proc:{}'.format(id), e)
    out_queue.put((None, True, id))
    if SHOW:
        cv2.destroyAllWindows()


def file_proc(file_queue):
    video_list_dict = OrderedDict()
    with open(video_list_file, 'r') as file:
        for line in file.readlines():
            line_info = line.strip().split(',')
            file_path = line_info[0]
            face_count = int(line_info[1])
            video_list_dict[file_path] = face_count

    _base_path = os.path.splitext(video_list_file)[0]
    title_file_path = _base_path + '.title.txt'
    with open(title_file_path, 'r') as file:
        title_str = file.readline()
        attr_item_count_str = file.readline()
        titles = title_str.strip().split(';')
        attr_item_counts = attr_item_count_str.strip().split(';')
        attr_item_counts = list(map(int, attr_item_counts))
        total_attr_item_count = sum(attr_item_counts)
        attr_count = len(titles)

    attr_item_offsets = [1]
    offset = 1
    for attr_item_count in attr_item_counts:
        offset += attr_item_count
        attr_item_offsets.append(offset)
    attrs_size = 4 + total_attr_item_count * 4
    # video_attributes_dict = {}
    data_file_path = _base_path + '.dat'
    with open(data_file_path, 'rb') as file:
        for video_name, face_count in video_list_dict.items():
            attributes_dict = {}
            for _ in range(face_count):
                dat = file.read(attrs_size)
                datas = struct.unpack('1i{}f'.format(total_attr_item_count), dat)
                frame_id = datas[0]
                attribute = {}
                for i, title in enumerate(titles):
                    start_offset = attr_item_offsets[i]
                    end_offset = attr_item_offsets[i+1]
                    attribute[title] = datas[start_offset:end_offset]
                attributes = attributes_dict.get(frame_id, None)
                if attributes is None:
                    attributes_dict[frame_id] = [attribute]
                else:
                    attributes.append(attribute)
            # video_attributes_dict[video_name] = attributes_dict
            while file_queue.qsize() > 100:
                time.sleep(0.01)
            file_queue.put((video_name, attributes_dict))
            # break

    # for video_name, attributes_dict in video_attributes_dict.items():
    #     while file_queue.qsize() > 100:
    #         time.sleep(0.01)
    #     file_queue.put((video_name, attributes_dict))

    for _ in range(PROCESS_COUNT):
        file_queue.put(None)


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    file_queue = multiprocessing.Queue()
    out_queue = multiprocessing.Queue()

    file_worker = multiprocessing.Process(target=file_proc, args=(file_queue,))

    file_worker.start()

    select_video_workers = []
    for i in range(PROCESS_COUNT):
        select_video_workers.append(multiprocessing.Process(target=crop_proc, args=(file_queue, out_queue, i,)))

    for i in range(PROCESS_COUNT):
        select_video_workers[i].start()

    file_count = 0
    try:
        finish_worker_count = 0
        while True:
            file_info = out_queue.get(block=True)
            if file_info is None:
                break
            file_name, finish, id = file_info
            if finish:
                print('Proc{} finish'.format(id, ))
                finish_worker_count += 1
                if PROCESS_COUNT <= finish_worker_count:
                    break
                continue
            file_count += 1
            print('Proc{}, File Count: {}, {}'.format(id, file_count, file_name))
    except Exception as e:
        print(e)

    for i in range(PROCESS_COUNT):
        select_video_workers[i].join()

    file_worker.join()

    os.system('chmod a+wr {} -Rf'.format(out_path))

    print('finish!!!!')


if __name__ == '__main__':
    main()
